// #include "CostVolumePlugin.h"
// #include "NvInfer.h"
// #include "cuda_runtime_api.h"
#include <cuda_fp16.h>
#include <vector>
#include <string>
#include <iostream>
#include <cuda_runtime_api.h>
#include "NvInferPlugin.h"
#include <cassert>
#include "CostVolumePlugin.hpp"
#include "CostVolume_kernel.hpp"
// #include "dsgn_kernel.h"

using namespace nvinfer1;
// using nvinfer1::plugin::CostVolumePlugin;
// using nvinfer1::plugin::CostVolumePluginCreator;

PluginFieldCollection CostVolumePluginCreator::mFC{};
std::vector<PluginField> CostVolumePluginCreator::mPluginAttributes;

REGISTER_TENSORRT_PLUGIN(CostVolumePluginCreator);

// Helper function for serializing plugin
template <typename T>
void writeToBuffer(char *&buffer, const T &val)
{
    *reinterpret_cast<T *>(buffer) = val;
    buffer += sizeof(T);
}

// Helper function for deserializing plugin
template <typename T>
T readFromBuffer(const char *&buffer)
{
    T val = *reinterpret_cast<const T *>(buffer);
    buffer += sizeof(T);
    return val;
}
namespace
{
    constexpr char const* COSTVOLUME_PLUGIN_NAME{"CostVolumePlugin"};
    constexpr char const* COSTVOLUME_VERSION{"1"};
}


// ##############################实现##########

CostVolumePlugin::CostVolumePlugin(int downsample){
    mDownsample = downsample;
}

CostVolumePlugin::CostVolumePlugin(void const* serialData, size_t length)
{
    // for deserialize
    const char* d = static_cast<const char *>(serialData);
    const char* a = d;
    mDownsample = readFromBuffer<int>(d);

    assert (d == (a + length));
}

int CostVolumePlugin::getNbOutputs() const noexcept
{
    return 1;
}

nvinfer1::DimsExprs CostVolumePlugin::getOutputDimensions(
    int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
{
    // Implement output dimensions calculation
    assert(outputIndex == 0);
    nvinfer1::DimsExprs output;
    output.nbDims = 5;
    output.d[0] = exprBuilder.constant(1);
    output.d[1] = exprBuilder.constant(64);
    output.d[2] = inputs[2].d[1];
    // output.d[2] = exprBuilder.constant(72);
    // output.d[3] = exprBuilder.constant(120);
    // output.d[4] = exprBuilder.constant(232);
    output.d[3] = inputs[0].d[2];
    output.d[4] = inputs[0].d[3];
    return output; // num_batch, sep * 2, max_disp, height, width

}
bool CostVolumePlugin::supportsFormatCombination(
    int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept
{
    // Implement format combination check
    //inOut : input, output
    assert(nbInputs == 3);
    assert(nbOutputs == 1);
    const PluginTensorDesc& in = inOut[pos];
    if (pos == 0)
    {
        return ((in.type == nvinfer1::DataType::kFLOAT) && in.format == TensorFormat::kLINEAR);
    }
    else if (pos == 1)
    {
        return ((in.type == nvinfer1::DataType::kFLOAT) && in.format == TensorFormat::kLINEAR);
    }
    else if (pos == 2)
    {
        return ((in.type == nvinfer1::DataType::kFLOAT) && in.format == TensorFormat::kLINEAR);
    }
    else if (pos == 3)
    {
        return (inOut[0].type == nvinfer1::DataType::kFLOAT && in.type == nvinfer1::DataType::kFLOAT && in.format == TensorFormat::kLINEAR);
    }
    // if (pos == 0)
    // {
    //     return ((in.type == nvinfer1::DataType::kFLOAT || in.type == nvinfer1::DataType::kHALF) && in.format == TensorFormat::kLINEAR);
    // }
    // else if (pos == 1)
    // {
    //     return ((in.type == nvinfer1::DataType::kFLOAT || in.type == nvinfer1::DataType::kHALF) && in.format == TensorFormat::kLINEAR);
    // }
    // else if (pos == 2)
    // {
    //     return ((in.type == nvinfer1::DataType::kFLOAT || in.type == nvinfer1::DataType::kHALF) && in.format == TensorFormat::kLINEAR);
    // }
    // else if (pos == 3)
    // {
    //     return (inOut[0].type == nvinfer1::DataType::kFLOAT && in.type == nvinfer1::DataType::kFLOAT && in.format == TensorFormat::kLINEAR)
    //         || (inOut[0].type == nvinfer1::DataType::kHALF  && in.type == nvinfer1::DataType::kHALF  && in.format == TensorFormat::kLINEAR);
    // }
    return false;
}
void CostVolumePlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs,
    const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept
{
    // Implement plugin configuration
    return;
}
size_t CostVolumePlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs,
    const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept
{
    // Implement workspace size calculation
    return 0;
}
nvinfer1::DataType CostVolumePlugin::getOutputDataType(
    int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
{
    // Implement output data type calculation
    assert(index == 0);
    return inputTypes[0];
}


const char* CostVolumePlugin::getPluginType() const noexcept
{
    return COSTVOLUME_PLUGIN_NAME;
}
const char* CostVolumePlugin::getPluginVersion() const noexcept
{
    return COSTVOLUME_VERSION;
}
int CostVolumePlugin::initialize() noexcept
{
    // Implement plugin initialization
    return 0;
}
void CostVolumePlugin::terminate() noexcept
{}

size_t CostVolumePlugin::getSerializationSize() const noexcept
{
    return sizeof(mDownsample);
}
void CostVolumePlugin::serialize(void *buffer) const noexcept
{
    char *p = static_cast<char *>(buffer);
    char *a = p;
    writeToBuffer(p, mDownsample);

    assert (p == (a + getSerializationSize()));
}
void CostVolumePlugin::destroy() noexcept
{
    // Implement plugin destruction
}
nvinfer1::IPluginV2DynamicExt* CostVolumePlugin::clone() const noexcept
{
    auto* plugin = new CostVolumePlugin(mDownsample);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

void CostVolumePlugin::setPluginNamespace(const char *pluginNamespace) noexcept
{
    // Implement plugin namespace setting
}
const char *CostVolumePlugin::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}
int CostVolumePlugin::enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
                                      const nvinfer1::PluginTensorDesc *outputDesc, 
                                      const void *const *inputs, void *const *outputs, void *workspace,
                                      cudaStream_t stream) noexcept
{
    // std::cout << "CostVolumePlugin enqueue" << std::endl;
    // Implement kernel launch
    int status = -1;
    int downsample = mDownsample;
    nvinfer1::DataType inputType = inputDesc[0].type;
    const nvinfer1::Dims inputDims = inputDesc[0].dims;
    const int numBatch = inputDims.d[0];
    const int channels = inputDims.d[1];
    const int height = inputDims.d[2] / downsample;
    const int width = inputDims.d[3] / downsample;
    const int max_disp = inputDesc[2].dims.d[1];
    if(inputType == nvinfer1::DataType::kFLOAT){
        auto left = static_cast<const float *>(inputs[0]);
        auto right = static_cast<const float *>(inputs[1]);
        auto shift = static_cast<const float *>(inputs[2]);
        auto cost = static_cast<float *>(outputs[0]);
        status = BuildCostVolumeKernelLaunch(left, right, shift,
                                            cost, numBatch, channels, height, width, max_disp,
                                            downsample, stream);
        assert (status == 0);
        return status;
    }
    // else if(inputType == nvinfer1::DataType::kHALF){
    //     auto left = static_cast<const half *>(inputs[0]);
    //     auto right = static_cast<const half *>(inputs[1]);
    //     auto shift = static_cast<const half *>(inputs[2]);
    //     auto psv_channels = static_cast<const int *>(inputs[3]);
    //     auto cost = static_cast<half *>(outputs[0]);
    //     status = BuildCostVolumeHalfKernelLaunch(left, right, shift,
    //                                                 cost, numBatch, channels, height, width, max_disp,
    //                                                 downsample, stream);
    //     assert (status == 0);
    //     return status;
    // }
        
    assert(status == 0);
    return status;
}


CostVolumePluginCreator::CostVolumePluginCreator(){
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(PluginField("downsample", nullptr, PluginFieldType::kINT32, 1));
    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}
const char* CostVolumePluginCreator::getPluginName() const noexcept
{
    return COSTVOLUME_PLUGIN_NAME;
}

const char* CostVolumePluginCreator::getPluginVersion() const noexcept
{
    return COSTVOLUME_VERSION;
}

const PluginFieldCollection* CostVolumePluginCreator::getFieldNames() noexcept
{
    return &mFC;
}


IPluginV2DynamicExt* CostVolumePluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
{   
    // std::cout << "CostVolumePluginCreator createPlugin" << std::endl;
    int downsample = 0;
    for (int i = 0; i < fc->nbFields; ++i)
    {
        const PluginField* field = fc->fields + i;
        if (field->data == nullptr)
        {
            continue;
        }
        std::string field_name(field->name);
        if (field_name.compare("downsample") == 0)
        {
            downsample = *static_cast<const int*>(field->data);
        }
    }
    CostVolumePlugin* plugin = new CostVolumePlugin(downsample);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

IPluginV2DynamicExt* CostVolumePluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call the destroy() method.
    CostVolumePlugin* plugin = new CostVolumePlugin(serialData, serialLength);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

void CostVolumePluginCreator::setPluginNamespace(const char* pluginNamespace) noexcept
{
    assert (pluginNamespace != nullptr);
    mNamespace = pluginNamespace;
}

const char* CostVolumePluginCreator::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

// REGISTER_TENSORRT_PLUGIN(CostVolumePluginCreator);